#include "LDPC_Matrix.h"

#include <fstream>

BaseGraph Load_LDPC_BG(void)
{
	BaseGraph bg;
	std::ifstream in("BG.dat");
	if (in.is_open() == true)
	{
		// Init BaseGraph
		bg.BG1_index = (uint32_t*)malloc(316 * sizeof(uint32_t));
		bg.BG1_Z1 = (uint32_t*)malloc(316 * sizeof(uint32_t));
		bg.BG1_Z2 = (uint32_t*)malloc(316 * sizeof(uint32_t));
		bg.BG1_Z3 = (uint32_t*)malloc(316 * sizeof(uint32_t));
		bg.BG1_Z4 = (uint32_t*)malloc(316 * sizeof(uint32_t));
		bg.BG1_Z5 = (uint32_t*)malloc(316 * sizeof(uint32_t));
		bg.BG1_Z6 = (uint32_t*)malloc(316 * sizeof(uint32_t));
		bg.BG1_Z7 = (uint32_t*)malloc(316 * sizeof(uint32_t));
		bg.BG1_Z8 = (uint32_t*)malloc(316 * sizeof(uint32_t));

		bg.BG2_index = (uint32_t*)malloc(197 * sizeof(uint32_t));
		bg.BG2_Z1 = (uint32_t*)malloc(197 * sizeof(uint32_t));
		bg.BG2_Z2 = (uint32_t*)malloc(197 * sizeof(uint32_t));
		bg.BG2_Z3 = (uint32_t*)malloc(197 * sizeof(uint32_t));
		bg.BG2_Z4 = (uint32_t*)malloc(197 * sizeof(uint32_t));
		bg.BG2_Z5 = (uint32_t*)malloc(197 * sizeof(uint32_t));
		bg.BG2_Z6 = (uint32_t*)malloc(197 * sizeof(uint32_t));
		bg.BG2_Z7 = (uint32_t*)malloc(197 * sizeof(uint32_t));
		bg.BG2_Z8 = (uint32_t*)malloc(197 * sizeof(uint32_t));

		bg.Zc = (uint32_t*)malloc(51 * sizeof(uint32_t));

		// Size of BG1 & BG2
		in >> bg.BG1_Row >> bg.BG1_Col >> bg.BG2_Row >> bg.BG2_Col;

		// Read BG1
		for (uint32_t i = 0; i < 316; ++i)
			in >> bg.BG1_index[i];
		for (uint32_t i = 0; i < 316; ++i)
			in >> bg.BG1_Z1[i];
		for (uint32_t i = 0; i < 316; ++i)
			in >> bg.BG1_Z2[i];
		for (uint32_t i = 0; i < 316; ++i)
			in >> bg.BG1_Z3[i];
		for (uint32_t i = 0; i < 316; ++i)
			in >> bg.BG1_Z4[i];
		for (uint32_t i = 0; i < 316; ++i)
			in >> bg.BG1_Z5[i];
		for (uint32_t i = 0; i < 316; ++i)
			in >> bg.BG1_Z6[i];
		for (uint32_t i = 0; i < 316; ++i)
			in >> bg.BG1_Z7[i];
		for (uint32_t i = 0; i < 316; ++i)
			in >> bg.BG1_Z8[i];

		// Read BG2
		for (uint32_t i = 0; i < 197; ++i)
			in >> bg.BG2_index[i];
		for (uint32_t i = 0; i < 197; ++i)
			in >> bg.BG2_Z1[i];
		for (uint32_t i = 0; i < 197; ++i)
			in >> bg.BG2_Z2[i];
		for (uint32_t i = 0; i < 197; ++i)
			in >> bg.BG2_Z3[i];
		for (uint32_t i = 0; i < 197; ++i)
			in >> bg.BG2_Z4[i];
		for (uint32_t i = 0; i < 197; ++i)
			in >> bg.BG2_Z5[i];
		for (uint32_t i = 0; i < 197; ++i)
			in >> bg.BG2_Z6[i];
		for (uint32_t i = 0; i < 197; ++i)
			in >> bg.BG2_Z7[i];
		for (uint32_t i = 0; i < 197; ++i)
			in >> bg.BG2_Z8[i];

		// Read Zc
		for (uint32_t i = 0; i < 51; ++i)
			in >> bg.Zc[i];

		in.close();
	}
	return bg;
}

void Free_LDPC_BG(BaseGraph& bg)
{
	free(bg.BG1_index);
	free(bg.BG1_Z1);
	free(bg.BG1_Z2);
	free(bg.BG1_Z3);
	free(bg.BG1_Z4);
	free(bg.BG1_Z5);
	free(bg.BG1_Z6);
	free(bg.BG1_Z7);
	free(bg.BG1_Z8);
	free(bg.BG2_index);
	free(bg.BG2_Z1);
	free(bg.BG2_Z2);
	free(bg.BG2_Z3);
	free(bg.BG2_Z4);
	free(bg.BG2_Z5);
	free(bg.BG2_Z6);
	free(bg.BG2_Z7);
	free(bg.BG2_Z8);
	free(bg.Zc);
}

void Generate_QC_Matrix(uint32_t Zc, uint32_t p, uint32_t* idx_r, uint32_t* idx_c)
{
	for (int i = 0; i < Zc; ++i)
		idx_r[i] = i;
	for (int i = 0; i < Zc; ++i)
		idx_c[i] = (p + i) % Zc;
}

bool Load_LDPC_Matrix(LDPC_Matrix& matrix, const int K, const float Rate)
{
	// Load Base Graph
	BaseGraph bg = Load_LDPC_BG();

	// Choose Base Graph, Compute Kb & Zc
	if (K <= 308 || (K <= 3840 && Rate <= 0.67) || Rate <= 0.25)
		matrix.BG_Choosen = 2;
	else
		matrix.BG_Choosen = 1;

	if (matrix.BG_Choosen == 1)
		matrix.Kb = 22;
	else
	{
		if (K > 640)
			matrix.Kb = 10;
		else if (K > 560)
			matrix.Kb = 9;
		else if (K > 192)
			matrix.Kb = 8;
		else
			matrix.Kb = 6;
	}
	for (int i = 0; i < 51; ++i)
	{
		if (matrix.Kb * bg.Zc[i] >= K)
		{
			matrix.Zc = bg.Zc[i];
			break;
		}
	}

	// Choose a_index and Point to Pointer
	uint32_t Base_Row, Base_Col;
	uint32_t* BG_index;
	uint32_t* BG_Z = nullptr;
	uint32_t NonZeroElemets = 0;
	if (matrix.BG_Choosen == 1)
	{
		Base_Row = bg.BG1_Row;
		Base_Col = bg.BG1_Col;
		BG_index = bg.BG1_index;
		NonZeroElemets = 316;
		if (matrix.Zc == 2 || matrix.Zc == 4 || matrix.Zc == 8 || matrix.Zc == 16 || matrix.Zc == 32 || matrix.Zc == 64 || matrix.Zc == 128 || matrix.Zc == 256)
		{
			BG_Z = bg.BG1_Z1; matrix.a_idx = 1;
		}
		else if (matrix.Zc == 3 || matrix.Zc == 6 || matrix.Zc == 12 || matrix.Zc == 24 || matrix.Zc == 48 || matrix.Zc == 96 || matrix.Zc == 192 || matrix.Zc == 384)
		{
			BG_Z = bg.BG1_Z2; matrix.a_idx = 2;
		}
		else if (matrix.Zc == 5 || matrix.Zc == 10 || matrix.Zc == 20 || matrix.Zc == 40 || matrix.Zc == 80 || matrix.Zc == 160 || matrix.Zc == 320)
		{
			BG_Z = bg.BG1_Z3; matrix.a_idx = 3;
		}
		else if (matrix.Zc == 7 || matrix.Zc == 14 || matrix.Zc == 28 || matrix.Zc == 56 || matrix.Zc == 112 || matrix.Zc == 224)
		{
			BG_Z = bg.BG1_Z4; matrix.a_idx = 4;
		}
		else if (matrix.Zc == 9 || matrix.Zc == 18 || matrix.Zc == 36 || matrix.Zc == 72 || matrix.Zc == 144 || matrix.Zc == 288)
		{
			BG_Z = bg.BG1_Z5; matrix.a_idx = 5;
		}
		else if (matrix.Zc == 11 || matrix.Zc == 22 || matrix.Zc == 44 || matrix.Zc == 88 || matrix.Zc == 176 || matrix.Zc == 352)
		{
			BG_Z = bg.BG1_Z6; matrix.a_idx = 6;
		}
		else if (matrix.Zc == 13 || matrix.Zc == 26 || matrix.Zc == 52 || matrix.Zc == 104 || matrix.Zc == 208)
		{
			BG_Z = bg.BG1_Z7; matrix.a_idx = 7;
		}
		else if (matrix.Zc == 15 || matrix.Zc == 30 || matrix.Zc == 60 || matrix.Zc == 120 || matrix.Zc == 240)
		{
			BG_Z = bg.BG1_Z8; matrix.a_idx = 8;
		}
	}
	else
	{
		Base_Row = bg.BG2_Row;
		Base_Col = bg.BG2_Col;
		BG_index = bg.BG2_index;
		NonZeroElemets = 197;
		if (matrix.Zc == 2 || matrix.Zc == 4 || matrix.Zc == 8 || matrix.Zc == 16 || matrix.Zc == 32 || matrix.Zc == 64 || matrix.Zc == 128 || matrix.Zc == 256)
		{
			BG_Z = bg.BG2_Z1; matrix.a_idx = 1;
		}
		else if (matrix.Zc == 3 || matrix.Zc == 6 || matrix.Zc == 12 || matrix.Zc == 24 || matrix.Zc == 48 || matrix.Zc == 96 || matrix.Zc == 192 || matrix.Zc == 384)
		{
			BG_Z = bg.BG2_Z2; matrix.a_idx = 2;
		}
		else if (matrix.Zc == 5 || matrix.Zc == 10 || matrix.Zc == 20 || matrix.Zc == 40 || matrix.Zc == 80 || matrix.Zc == 160 || matrix.Zc == 320)
		{
			BG_Z = bg.BG2_Z3; matrix.a_idx = 3;
		}
		else if (matrix.Zc == 7 || matrix.Zc == 14 || matrix.Zc == 28 || matrix.Zc == 56 || matrix.Zc == 112 || matrix.Zc == 224)
		{
			BG_Z = bg.BG2_Z4; matrix.a_idx = 4;
		}
		else if (matrix.Zc == 9 || matrix.Zc == 18 || matrix.Zc == 36 || matrix.Zc == 72 || matrix.Zc == 144 || matrix.Zc == 288)
		{
			BG_Z = bg.BG2_Z5; matrix.a_idx = 5;
		}
		else if (matrix.Zc == 11 || matrix.Zc == 22 || matrix.Zc == 44 || matrix.Zc == 88 || matrix.Zc == 176 || matrix.Zc == 352)
		{
			BG_Z = bg.BG2_Z6; matrix.a_idx = 6;
		}
		else if (matrix.Zc == 13 || matrix.Zc == 26 || matrix.Zc == 52 || matrix.Zc == 104 || matrix.Zc == 208)
		{
			BG_Z = bg.BG2_Z7; matrix.a_idx = 7;
		}
		else if (matrix.Zc == 15 || matrix.Zc == 30 || matrix.Zc == 60 || matrix.Zc == 120 || matrix.Zc == 240)
		{
			BG_Z = bg.BG2_Z8; matrix.a_idx = 8;
		}
	}


	// Position of Nonzero Elements of Base Graph & Check Matrix
	matrix.H_Complete_NonzeroElements = NonZeroElemets * matrix.Zc;
	uint32_t* BG_idx_r = new uint32_t[NonZeroElemets];
	uint32_t* BG_idx_c = new uint32_t[NonZeroElemets];
	uint32_t* Z_idx_r = new uint32_t[matrix.Zc];
	uint32_t* Z_idx_c = new uint32_t[matrix.Zc];
	matrix.H_Complete_Row_Index = new uint32_t[matrix.H_Complete_NonzeroElements];
	matrix.H_Complete_Col_Index = new uint32_t[matrix.H_Complete_NonzeroElements];



	for (int i = 0; i < NonZeroElemets; ++i)
	{
		BG_idx_r[i] = (BG_index[i] - 1) % Base_Row;
		BG_idx_c[i] = (BG_index[i] - 1) / Base_Row;
		Generate_QC_Matrix(matrix.Zc, BG_Z[i], Z_idx_r, Z_idx_c);
		for (int k = 0; k < matrix.Zc; ++k)
		{
			matrix.H_Complete_Row_Index[i * matrix.Zc + k] = Z_idx_r[k] + BG_idx_r[i] * matrix.Zc;
			matrix.H_Complete_Col_Index[i * matrix.Zc + k] = Z_idx_c[k] + BG_idx_c[i] * matrix.Zc;
		}
	}


	// Cut Check Matrix
	matrix.nbrOfInfoBits = K;
	matrix.nbrOfCheckBits = ((round(K / Rate) - K + 2 * matrix.Zc) > matrix.Zc * 4) ? (round(K / Rate) - K + 2 * matrix.Zc) : matrix.Zc * 4;
	matrix.beginOfCheckBit = (matrix.BG_Choosen == 1) ? 22 * matrix.Zc : 10 * matrix.Zc;
	matrix.nbrOfRow = Base_Row * matrix.Zc;
	matrix.nbrOfCol = Base_Col * matrix.Zc;

	int nonzeroelements = 0;
	for (int i = 0; i < NonZeroElemets * matrix.Zc; ++i)
	{
		if ((matrix.H_Complete_Col_Index[i] < K // Info Bit
			|| ((matrix.H_Complete_Col_Index[i] >= matrix.beginOfCheckBit) && matrix.H_Complete_Col_Index[i] < matrix.beginOfCheckBit + matrix.nbrOfCheckBits)) // Check Bit
			&& (matrix.H_Complete_Row_Index[i] < matrix.nbrOfCheckBits)) // Row
			++nonzeroelements;
	}

	// Save Check Matrix Nonzero Element Index
	matrix.H_Row_Index = new uint32_t[nonzeroelements];
	matrix.H_Col_Index = new uint32_t[nonzeroelements];
	matrix.H_NonzeroElements = nonzeroelements;
	int index = 0;
	for (int i = 0; i < NonZeroElemets * matrix.Zc; ++i)
	{
		if ((matrix.H_Complete_Col_Index[i] < K // Info Bit
			|| ((matrix.H_Complete_Col_Index[i] >= matrix.beginOfCheckBit) && matrix.H_Complete_Col_Index[i] < matrix.beginOfCheckBit + matrix.nbrOfCheckBits)) // Check Bit
			&& (matrix.H_Complete_Row_Index[i] < matrix.nbrOfCheckBits)) // Row
		{
			matrix.H_Row_Index[index] = matrix.H_Complete_Row_Index[i];
			matrix.H_Col_Index[index] = matrix.H_Complete_Col_Index[i];
			++index;
		}
	}
	matrix.nbrOfCheckBits = round(K / Rate - K + 2 * matrix.Zc);
	matrix.Rate = matrix.nbrOfInfoBits / (double)(matrix.nbrOfInfoBits + matrix.nbrOfCheckBits - 2 * matrix.Zc);

	// Cut Base Matrix
	nonzeroelements = 0;
	for (int i = 0; i < matrix.H_NonzeroElements; ++i)
	{
		if (matrix.H_Col_Index[i] < K
			&& matrix.H_Row_Index[i] < 4 * matrix.Zc)
			++nonzeroelements;
	}

	// Save Base Matrix Nonzero Element Index
	matrix.H_Base_Row_Index = new uint32_t[nonzeroelements];
	matrix.H_Base_Col_Index = new uint32_t[nonzeroelements];
	matrix.H_Base_NonzeroElements = nonzeroelements;
	index = 0;
	for (int i = 0; i < matrix.H_NonzeroElements; ++i)
	{
		if (matrix.H_Col_Index[i] < K
			&& matrix.H_Row_Index[i] < 4 * matrix.Zc)
		{
			matrix.H_Base_Row_Index[index] = matrix.H_Row_Index[i];
			matrix.H_Base_Col_Index[index] = matrix.H_Col_Index[i];
			++index;
		}
	}

	// Cut Expand Matrix
	nonzeroelements = 0;
	for (int i = 0; i < matrix.H_Complete_NonzeroElements; ++i)
	{
		if (matrix.H_Complete_Col_Index[i] < matrix.beginOfCheckBit + 4 * matrix.Zc
			&& matrix.H_Complete_Row_Index[i] >= 4 * matrix.Zc)
			++nonzeroelements;
	}

	// Save Expand Matrix Nonzero Element Index
	matrix.H_Expand_NonzeroElements = nonzeroelements;
	if (nonzeroelements != 0)
	{
		matrix.H_Expand_Row_Index = new uint32_t[nonzeroelements];
		matrix.H_Expand_Col_Index = new uint32_t[nonzeroelements];
		index = 0;
		for (int i = 0; i < matrix.H_Complete_NonzeroElements; ++i)
		{
			if (matrix.H_Complete_Col_Index[i] < matrix.beginOfCheckBit + 4 * matrix.Zc
				&& matrix.H_Complete_Row_Index[i] >= 4 * matrix.Zc)
			{
				matrix.H_Expand_Row_Index[index] = matrix.H_Complete_Row_Index[i];
				matrix.H_Expand_Col_Index[index] = matrix.H_Complete_Col_Index[i];
				++index;
			}
		}
	}
	else
	{
		matrix.H_Expand_Row_Index = nullptr;
		matrix.H_Expand_Col_Index = nullptr;
	}

	// Sort
	uint32_t temp;
	for (int i = 0; i < matrix.H_NonzeroElements; ++i)
		for (int j = matrix.H_NonzeroElements - 1; j > i; --j)
			if (matrix.H_Row_Index[j] < matrix.H_Row_Index[j - 1])
			{
				temp = matrix.H_Row_Index[j];
				matrix.H_Row_Index[j] = matrix.H_Row_Index[j - 1];
				matrix.H_Row_Index[j - 1] = temp;

				temp = matrix.H_Col_Index[j];
				matrix.H_Col_Index[j] = matrix.H_Col_Index[j - 1];
				matrix.H_Col_Index[j - 1] = temp;
			}

	// Generate Connection Relationship
	uint32_t sumidxq = 0, sumidxr = 0;
	for (uint32_t i = 0; i < matrix.H_NonzeroElements; i++)
	{
		for (uint32_t j = 0; j < matrix.H_NonzeroElements; j++)
			if (matrix.H_Row_Index[j] == matrix.H_Row_Index[i] && matrix.H_Col_Index[j] != matrix.H_Col_Index[i])
				sumidxq++;
		for (uint32_t j = 0; j < matrix.H_NonzeroElements; j++)
			if (matrix.H_Col_Index[j] == matrix.H_Col_Index[i] && matrix.H_Row_Index[j] != matrix.H_Row_Index[i])
				sumidxr++;
	}
	matrix.IdxConnectQij = new uint32_t[sumidxq];
	matrix.IdxConnectRji = new uint32_t[sumidxr];
	matrix.IdxConnectQij_idx = new uint32_t[(__int64)matrix.H_NonzeroElements * 2];
	matrix.IdxConnectRji_idx = new uint32_t[(__int64)matrix.H_NonzeroElements * 2];
	sumidxq = 0;
	sumidxr = 0;
	for (uint32_t i = 0; i < matrix.H_NonzeroElements; i++)
	{
		matrix.IdxConnectQij_idx[2 * i] = sumidxq;
		for (uint32_t j = 0; j < matrix.H_NonzeroElements; j++)
			if (matrix.H_Row_Index[j] == matrix.H_Row_Index[i] && matrix.H_Col_Index[j] != matrix.H_Col_Index[i])
				matrix.IdxConnectQij[sumidxq++] = j;
		matrix.IdxConnectQij_idx[2 * i + 1] = sumidxq;

		matrix.IdxConnectRji_idx[2 * i] = sumidxr;
		for (uint32_t j = 0; j < matrix.H_NonzeroElements; j++)
			if (matrix.H_Col_Index[j] == matrix.H_Col_Index[i] && matrix.H_Row_Index[j] != matrix.H_Row_Index[i])
				matrix.IdxConnectRji[sumidxr++] = j;
		matrix.IdxConnectRji_idx[2 * i + 1] = sumidxr;
	}

	// Free Space
	Free_LDPC_BG(bg);
	delete[]BG_idx_r;
	delete[]BG_idx_c;
	delete[]Z_idx_r;
	delete[]Z_idx_c;

	return true;
}

void Free_LDPC_Matrix(LDPC_Matrix& matrix)
{
	delete[]matrix.H_Row_Index;
	delete[]matrix.H_Col_Index;
	delete[]matrix.H_Base_Row_Index;
	delete[]matrix.H_Base_Col_Index;
	delete[]matrix.H_Complete_Row_Index;
	delete[]matrix.H_Complete_Col_Index;
	if (matrix.H_Expand_Row_Index != nullptr)
	{
		delete[]matrix.H_Expand_Row_Index;
		delete[]matrix.H_Expand_Col_Index;
	}


	delete[]matrix.IdxConnectQij;
	delete[]matrix.IdxConnectRji;
	delete[]matrix.IdxConnectQij_idx;
	delete[]matrix.IdxConnectRji_idx;
}